#!/usr/bin/env python
# -*- coding: utf8 -*-
# title       :
# description :
# author      :'ShenMeng'


import pymel.core as pmel
import maya.cmds as cmds
import maya.mel as mel
import os
import re
import yaml
from verexport.core import VerFbxExport, VerAbcExport, VerBakeAnimation


class ExportManager():
    def __init__(self, config):
        self._config = config
        self._initData()
        try:
            cmds.loadPlugin('AbcExport')
        except:pass
        try:
            cmds.loadPlugin('fbxmaya')
        except:pass

    def _initData(self):
        if os.path.isfile(self._config):
            configFile = self._config
        else:
            config_path = os.getenv("EXPORT_MANAGER_CONFIG_PATH")
            if config_path:
                configFile = os.path.join(config_path, self._config)
            else:
                configFile = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'config', self._config).replace('\\', '/')
        if not os.path.isfile(configFile):
            raise ValueError("config file is not exists!\n %s " % configFile)
        with open(configFile) as f:
            self._data = yaml.load(f)
            
    def getScenePice(self):
        fileName = os.path.basename(pmel.sceneName())
        sceneName = os.path.splitext(fileName)[0]
        namePart = []
        outPath = os.getenv("EXPORT_MANAGER_OUT_PATH")
        if outPath and os.path.isdir(outPath):
            outPath = outPath.replace('\\', '/')
            namePart.append(outPath)
        else:
            namePart.append(self._data['outPath'])
        namePart.extend(sceneName.split('_'))
        return [sceneName, namePart]
        
    def main(self):
        fbxExportData = {}
        sceneName, namePart = self.getScenePice()
        fbxGroups = self._data['fbx']['groups']
        abcGroups = self._data['abc']['groups']
        #>> fbx & abc
        refs = pmel.listReferences()
        for ref in refs:
        
            if not ref.isLoaded():
                continue
            assetFolder = os.path.dirname(ref.path)
            assetType = os.path.basename(os.path.dirname(assetFolder))
            if assetType not in self._data['exportType'].keys():
                continue
            assetName = os.path.basename(assetFolder)
            subData = ref.subReferences()
            if subData:
                for key, subref in subData.iteritems():
                    if not subref.path.endswith('_ea.mb'):
                        ue_file = subref.path.replace('.mb', self._data['reference']['replace'])
                        if os.path.isfile(ue_file):
                            subref.replaceWith(ue_file)
            else:
                if not ref.path.endswith('_ea.mb'):
                    if not ref.path.endswith(self._data['reference']['replace']):
                        ue_file = ref.path.replace('.mb', self._data['reference']['replace'])
                        if os.path.isfile(ue_file):
                            ref.replaceWith(ue_file)

            #
            refNamespace = ref.namespace
            mat = re.match('^{assetName}([a-z_]+)(|[0-9][0-9]|[0-9])$'.format(assetName=assetName), refNamespace)
            if mat:
                instanceName = refNamespace.replace(mat.group(1), '')
            else:
                instanceName = refNamespace
            #
            #ref.importContents()
            rootNode = self.getRootNode(refNamespace)
            #cmds.namespace(mv=[refNamespace, ':'], f=1)

            fbxResultGrp = []
            abcResultGrp = []
            fbxExportData[instanceName] = {'asset':assetName, 'node':[], 'ref':ref, 'type':assetType, 'root':rootNode}
            for node in rootNode.getChildren():
                for grp in abcGroups: 
                    if node.name().endswith(grp):
                        abcResultGrp.append(node.name())
                        ####
                        if ref.path.endswith('_ea.mb'):
                            all_mesh = cmds.ls(node.name(), dag=1, type="mesh")
                            if all_mesh:
                                for mesh in all_mesh:
                                    self.smoothNode(mesh)
                        ####
                        abcCombine = pmel.polyUnite(node.name(), ch=1, mergeUVSets=1, name=instanceName)
                        pmel.parent(abcCombine, node)
                        framePiceList = self.getFramePice(self._data['frameStep'])
                        for index, frameRange in enumerate(framePiceList):
                            abcNamePart = namePart[:]
                            abcNamePart.append(self._data['exportType'][assetType])
                            abcFileName = '_'.join([sceneName, instanceName, '%02d' % (index+1), '{}.abc'.format(self._data['abc']['postName'])])
                            abcNamePart.append(abcFileName)
                            abcOutPath = os.path.join(*abcNamePart).replace('\\', '/')
                            abcDataPath = abcOutPath.replace('.abc', '.txt')
                            self.validatePath(abcOutPath)
                            self.exportAbc([abcCombine[0].name()], abcOutPath, frameRange)
                            self.validatePath(abcDataPath)
                            self.write(abcDataPath, frameRange)
                        break
                if node.name() in abcResultGrp:
                    continue
                for grp in fbxGroups:
                    if node.name().endswith(grp):
                        fbxExportData[instanceName]['node'].append(node)
                        break


        ###---back---###
        # TODO
        #execfile("H:/0_RD/Ver_AniTools/SSLee_AniTools/py/bake_anm_hook.py")
        self.bakeAnim(self.getBakeNodes())
        
        for instanceName, data in fbxExportData.iteritems():
            print '%s >>' % instanceName
            ref = data['ref']
            rootNode = data['root']
            exportNodes = data['node']
            refNamespace = ref.namespace
            subData = ref.subReferences()
            ref.importContents()
            if subData:
                for key, subref in subData.iteritems():
                    subref.importContents()
            cmds.namespace(mv=[refNamespace, ':'], f=1)
            assetName = data['asset']
            assetType = data['type']
            fbxNamePart = namePart[:]
            fbxNamePart.append(self._data['exportType'][assetType])
            fbxFileName = '_'.join([sceneName, instanceName, '{}.fbx'.format(self._data['fbx']['postName'])])
            fbxNamePart.append(fbxFileName)
            fbxOutPath = os.path.join(*fbxNamePart).replace('\\', '/')
            fbxAssetDataPath = fbxOutPath.replace('{}.fbx'.format(self._data['fbx']['postName']), '{}.txt'.format(self._data['fbx']['postName']))
            # 2.export
            self.validatePath(fbxOutPath)
            self.exportFbx(exportNodes, fbxOutPath)
            # 3.write
            self.validatePath(fbxAssetDataPath)
            self.write(fbxAssetDataPath, assetName)
            rootNode.rename(refNamespace+':'+rootNode.name())
        # >> camera & scene frame
        #
        frameData = self.getFrameRange()
        allCameras = self.getCameras()

        self.bakeAnim(allCameras, frameData)
        #cmds.bakeResults(allCameras, simulation=1, t=frameData, controlPoints=1)
        for cam in allCameras:
            camDataPath = '{}/{}'.format(os.path.join(*namePart), '{}_{}.fbx'.format(sceneName, cam)).replace('\\', '/')
            self.exportFbx(cam, camDataPath)
        #
        sceneFrameDataPath = '{}/{}'.format(os.path.join(*namePart), '{}.txt'.format(sceneName)).replace('\\', '/')
        self.validatePath(sceneFrameDataPath)
        self.write(sceneFrameDataPath, frameData)
    @staticmethod
    def smoothNode(node):
        cmds.polySmooth(node, mth=0, sdt=2, ovb=1, ofb=3, ofc=0, ost=1, ocr=0, dv=1, 
                        bnr=1, c=1, kb=1, ksb=1, khe=0, kt=1, kmb=1, suv=1, peh=0, sl=1, dpe=1, ps=0.1, ro=1, ch=1)
    @staticmethod
    def exportAbc(nodes, outPath, frameRange):
        exporter = VerAbcExport(filePath=outPath)
        exporter.exportObjects = nodes
        exporter.frameRange = frameRange
        exporter.main()
        return
        cmds.AbcExport(j="-frameRange {} {} -uvWrite -writeFaceSets -worldSpace -writeVisibility -dataFormat ogawa -root {} -file {}".format(frameRange[0], frameRange[1], nodes[0], outPath))
        
    @staticmethod
    def exportFbx(nodes, outPath):
        exporter = VerFbxExport(filePath=outPath)
        exporter.exportObjects = nodes
        exporter.main()
        return
        cmdstr = """
        FBXProperty Export|IncludeGrp|Geometry|SmoothingGroups   -v true;
        FBXProperty Export|IncludeGrp|Animation -v true;
        FBXProperty Export|IncludeGrp|Animation|BakeComplexAnimation  -v true;
        FBXProperty Export|IncludeGrp|Animation|Deformation   -v true;  
        FBXProperty Export|IncludeGrp|Animation|Deformation|Skins  -v true;   
        FBXProperty Export|IncludeGrp|Animation|Deformation|Shape   -v true;   
        FBXProperty Export|IncludeGrp|Animation|ConstraintsGrp|Constraint   -v true;
        FBXProperty Export|IncludeGrp|CameraGrp|Camera -v true;
        FBXProperty Export|AdvOptGrp|AxisConvGrp|UpAxis -v "Y";
        """
        mel.eval(cmdstr)
        pmel.select(cl=1)
        pmel.select(nodes)
        cmds.file(outPath, force=1, options="v=0;", typ="FBX export", pr=1, es=1)
        
    @staticmethod
    def validatePath(filePath):
        if not os.path.isdir(os.path.dirname(filePath)):
            os.makedirs(os.path.dirname(filePath))
        #elif os.path.isfile(filePath):
        #    os.remove(filePath)
    
    @staticmethod
    def getFrameRange():
        st = cmds.playbackOptions(q=1, ast=1)
        et = cmds.playbackOptions(q=1, aet=1)
        return (st, et)
    
    @staticmethod
    def write(filePath, data):
        with open(filePath, 'w') as f:
            f.write(str(data))

    def getCameras(self):
        result = []
        regex = self._data['camera'].get('regex')
        if not regex:
            result.append(self._data['camera']['name'])
        else:
            for cam in pmel.ls(type='camera'):
                camName = cam.getParent().name()
                if re.match(regex, camName):
                    result.append(camName)
        return result
        
    def getFramePice(self, step):
        result = []
        st, et = self.getFrameRange()
        _temp = st
        while (_temp + step <= et and (et-(_temp + step))>30):
            result.append([_temp, _temp + step])
            _temp = _temp + step + 1
        result.append([_temp, et])
        return result
        
    @staticmethod
    def getRootNode(namespace):
        for node in pmel.ls(assemblies=1):
            if node.namespace() == namespace+':':
                return node
                
    def getBakeNodes(self):
        blendshapeList = self.getFaceBlendshape()
        jointList = cmds.ls(type='joint')
        nodeList = []
        nodeList.extend(blendshapeList)
        nodeList.extend(jointList)
        return nodeList
        
    @staticmethod
    def bakeAnim(nodes, frameRange=None):
        baker = VerBakeAnimation()
        baker.bakeObjects = nodes
        if frameRange:
            baker.frameRange = frameRange
        baker.main()

        
    @staticmethod
    def getFaceBlendshape():
        def getParent(node):
            parentNode = node.getParent()
            
            if not parentNode:
                return False
            else:
                if parentNode.endswith('face_grp'):
                    print parentNode
                    return True
                else:
                    return getParent(parentNode)
        allBlendShape = cmds.ls(type='blendShape')
        resultBlendShape = []
        for bs in allBlendShape:
            srcObj = cmds.blendShape(bs, q=1, g=1)[0]
            srcNode = pmel.PyNode(srcObj)
            if getParent(srcNode):
                resultBlendShape.append(bs)
        
        return list(set(allBlendShape)-set(resultBlendShape))

if __name__ == "__main__":
    em = ExportManager(config='deformationcar.yml')
    em.main()
